#include <iostream>
#include "../common/util.h"
#include "mem.h"

HeapAllocator::HeapAllocator(const char* name, char* base, U32 free_size)
{
    mBase = (U32)base;
    mBrk = mBase;
    mFreeMemory = free_size;
}

void* HeapAllocator::allocate(U32 t, U32 align)
{
    ASSERT(align != 0);

    // unaligned raw address
    U32 rawAddress;

    // Determine total amount of memory to allocate
    U32 expandedSize_bytes = t + align;

    ASSERT(expandedSize_bytes < mFreeMemory);

    // Allocate an unaligned block & convert address to U32
    // and adjust the heap top
    rawAddress = mBrk;
    mBrk += expandedSize_bytes;

    // assert there is enough space to allocate a new class on the static heap
    ASSERT(mFreeMemory > expandedSize_bytes);
    mFreeMemory -= expandedSize_bytes;


    // Calculate the adjustment by masking off the lower
    // bits of the address, to determine how "misaligned"
    // it is
    U32 mask = align - 1;
    U32 misalign = (rawAddress & mask);
    U32 adjustment = align - misalign;

    // Calculate the adjusted address, and return as
    // a pointer
    U32 alignedAddress = rawAddress + adjustment;

    // store the adjustment in the four bytes immediately
    // before the adjusted address that we are returning
    U32* pAdjustment = (U32*)(alignedAddress - 4);
    *pAdjustment = adjustment;

    return (void*)alignedAddress;
}

void HeapAllocator::deallocate(void* p)
{
    ASSERT((U32)p < mBrk);

    // read the alignment of the memory, which resides at
    // the byte immediately before p, and then calculate
    // the raw address allocated
    U32 alignedAddress = (U32)p;
    U32* pAdjustment = (U32*)(alignedAddress - 4);
    U32 adjustment = *pAdjustment;

    U32 rawAddress = alignedAddress - adjustment;

    // adjust free memory count and heap top
    mFreeMemory += mBrk - rawAddress;
    mBrk = rawAddress;

}

PoolAllocator::PoolAllocator(const char* name, const size_t& totalMemorySize, const size_t& chunkSize)
{
    mpFirstChunk = 0;
    mpLastChunk = 0;

    mpCurChunk = 0;

    mTotalMemorySize = 0;
    mUsedMemorySize = 0;
    mFreeMemorySize = 0;

    // set chunk size
    mChunkSize = chunkSize;

    AllocateMemory(totalMemorySize);

}

void PoolAllocator::AllocateMemory(const size_t& memorysize)
{
    // malloc enough memory for the chunks used
    U32 chunkCount = CalculateNeededChunks(memorysize);
    U8* pMemory = (U8*)malloc(chunkCount * mChunkSize);
    ASSERT(pMemory != 0);

    // malloc enough memory for data
    MemoryChunk* pChunks = new MemoryChunk[chunkCount];
    ASSERT(pChunks != 0);

    // update allocator information
    mTotalMemorySize += chunkCount * mChunkSize;
    mFreeMemorySize += chunkCount * mChunkSize;

    // link data to chunks
    LinkChunkToData(pMemory, pChunks, chunkCount);
}

void PoolAllocator::LinkChunkToData(U8* pMemory, MemoryChunk* pChunks, U32 chunkCount)
{
    U32 offsetMemory;

    // set first chunk and last chunk properly
    mpFirstChunk = &pChunks[0];
    mpLastChunk = &pChunks[chunkCount-1];

    for(U32 i = 0; i < chunkCount; i++)
    {
        offsetMemory = i * (U32)mChunkSize;

        // set each chunk's data to proper data addresses
        pChunks[i].data = &pMemory[offsetMemory];

        // dataSize indicates how many memory is available
        // starting from this chunk. At begining, because
        // all the chunks are available, the dataSize is the
        // size of memory from this chunk to the end.
        pChunks[i].dataSize = (U32)mTotalMemorySize - offsetMemory;
    }

}

void* PoolAllocator::allocate(U32 t, U32 align)
{
    ASSERT(align != 0);

    // unaligned raw address
    U32 rawAddress;

    // Determine total amount of memory to allocate
    U32 expandedSize_bytes = t + align;

    // Allocate an unaligned block
    rawAddress = (U32)AllocateUnalignedMemory(expandedSize_bytes);

    // Calculate the adjustment by masking off the lower
    // bits of the address, to determine how "misaligned"
    // it is
    U32 mask = align - 1;
    U32 misalign = (rawAddress & mask);
    U32 adjustment = align - misalign;

    // Calculate the adjusted address, and return as
    // a pointer
    U32 alignedAddress = rawAddress + adjustment;

    // store the adjustment in the four bytes immediately
    // before the adjusted address that we are returning
    U32* pAdjustment = (U32*)(alignedAddress - 4);
    *pAdjustment = adjustment;

    return (void*)alignedAddress;
}

void* PoolAllocator::AllocateUnalignedMemory(U32 size)
{
    MemoryChunk* pNewChunk = 0;

    // size cannot be 0, because
    // if it's 0, nothing will be
    // allocated, FindChunkToHoldMemory
    // can be undefined.
    ASSERT(size>0);
    U32 chunkCount = CalculateNeededChunks(size);
    U32 requestedMemory = chunkCount * mChunkSize;

    #ifdef _DEBUG
        std::cerr<<"searching for memory of size "<<requestedMemory<<std::endl;
    #endif

    // find a chunk that can hold the requested memory
    // if it fails, a NULL is returned.
    pNewChunk = FindChunkToHoldMemory(requestedMemory);
    ASSERT(pNewChunk != 0);

    // update allocator information
    mUsedMemorySize += requestedMemory;
    mFreeMemorySize -= requestedMemory;

    // properly set the value of head and tail chunk of allocation
    // 1. set their data size to be 0
    pNewChunk->dataSize = 0;
    pNewChunk[chunkCount-1].dataSize = 0;

    // 2. set their used data as request memory
    pNewChunk->usedSize = requestedMemory;
    pNewChunk[chunkCount-1].usedSize = requestedMemory;
    
    return (void*)pNewChunk->data;
}

MemoryChunk* PoolAllocator::FindChunkToHoldMemory(U32 memorySize)
{
    MemoryChunk* pMemoryChunk;

    // search for qualified chunk
    pMemoryChunk = mpCurChunk;
    U32 chunksToSkip = 0;
    do
    {
        if(pMemoryChunk == mpLastChunk || pMemoryChunk == 0)
        {
            // End of List reached : Start over from the beginning
            pMemoryChunk = mpFirstChunk ;
        }
        else
        {
            // proceed
            pMemoryChunk++;
        }

        if(pMemoryChunk != 0)
        {

            if(pMemoryChunk->dataSize >= memorySize && pMemoryChunk->usedSize == 0)
            {
                // a suitable chunk is found, set cur chunk to its tail chunk of this
                // memory segment.
                mpCurChunk = pMemoryChunk+CalculateNeededChunks(memorySize) - 1;
                return pMemoryChunk;
            }
            
            if(pMemoryChunk->usedSize != 0)
            {
                // this chunk is already allocated. then skip to the tail of this
                // memory segment
                chunksToSkip = CalculateNeededChunks(pMemoryChunk->usedSize);
                pMemoryChunk += chunksToSkip - 1;
            }
	}

    }while(pMemoryChunk != mpCurChunk);
    return 0;
}

void PoolAllocator::deallocate(void* p)
{
    // calculate the raw address. the alignment is store at
    // the place immediately before the aligned address, so
    // first read the alignment and then shift the address
    // to get raw address
    U32 alignedAddress = (U32)p;
    U32* pAdjustment = (U32*)(alignedAddress - 4);
    U32 adjustment = *pAdjustment;

    U32 rawAddress = alignedAddress - adjustment;
    ASSERT(rawAddress >= (U32)mpFirstChunk->data);
    ASSERT(rawAddress < ((U32)mpFirstChunk->data) + mTotalMemorySize);
    U32 index;

    // calculate the index of chunk
    index = (rawAddress - (U32)mpFirstChunk->data) / mChunkSize;

    // calcualte the begin chunk and the end chunk of
    // this memory segment
    MemoryChunk* pBeginChunk = &mpFirstChunk[index];
    MemoryChunk* pCurChunk = pBeginChunk;
    U32 chunkCount = CalculateNeededChunks(pCurChunk->usedSize);
    MemoryChunk* pEndChunk = &pCurChunk[chunkCount - 1];

    // reset their usedSize
    pEndChunk->usedSize = 0;
    pBeginChunk->usedSize = 0;

    // adjust the data size information of chunks between
    // begin chunk and end chunk.
    for(U32 i = 0;i<chunkCount;i++)
    {
        ASSERT(pCurChunk != 0);
        if(pEndChunk != mpLastChunk)
        {
            pCurChunk->dataSize = (pEndChunk + 1)->dataSize + mChunkSize*(chunkCount - i);
        }
        else
        {
            pCurChunk->dataSize = mChunkSize*(chunkCount - i);
        }
        pCurChunk++;
    }

    mUsedMemorySize -= chunkCount * mChunkSize;
    mFreeMemorySize += chunkCount * mChunkSize;

    if(pBeginChunk != mpFirstChunk)
    {
        // update the information of free chunks before deallocated chunk
        for(U32 i = 1; i<=index && mpFirstChunk[index-i].dataSize != 0 ; i++)
        {
            mpFirstChunk[index-i].dataSize = pBeginChunk->dataSize + i*mChunkSize;
        }
    }
}


